# -*- coding: utf-8 -*-
import numpy as np
import scipy
import matplotlib.pyplot as plt
import matplotlib


class RoughProfile(object):
    def __init__(self):
        self.step = 0
        self.altitude = []
        self.illegal_positions = []
        self.parameters = {}    # 存储表面轮廓表征参数的字典

    def show_profile(self):
        """
        绘制轮廓

        """
        matplotlib.rcParams['font.family'] = ['STSong']
        fig, ax = plt.subplots(1, 1)
        num_points = np.size(self.altitude)
        x_values = np.arange(0, num_points*self.step, self.step)
        ax.set_title("表面轮廓")
        ax.plot(x_values, self.altitude)
        fig.show()


class MeasuredProfile(RoughProfile):
    def __init__(self):
        super(MeasuredProfile, self).__init__()
        self.temporary_altitude = None

    def show_temporary_profile(self):
        pass


class FractalProfile(RoughProfile):
    def __init__(self, num_points, dimension, step, Rq, stable=True):
        super(FractalProfile, self).__init__()

        self.dimension = dimension
        self.num_points = num_points
        self.step = step
        self.Rq = Rq
        self.stable = stable

    def get_fractal_profile_DFT(self):
        """
        基于离散傅里叶逆变换生成分形轮廓

        :param num_points: 轮廓的采样点个数，其值应该是2的幂次方
        :param dimension: 分形维数，其值应该在区间(1, 2)之间
        :param Rq: 采样点高度均方根偏差的期望值
        :param stable: 是否生成平稳分形轮廓，如果是平稳分形轮廓，那么生成轮廓的Rq值等于输入的Rq值
        :return: 分形轮廓
        """

        # 判断采样点数（num_points）是否为2的幂次方，如果不是，则中断执行
        import math
        assert math.log(self.num_points, 2).is_integer(), "输入的采样点数不是2的幂次方，程序退出"
        # 判断分形维数是否在区间(1, 2)，如果不是，则中断执行
        assert 1 < self.dimension < 2, "分形维数必须在1和2之间，程序退出"

        tmp_1 = 5 - 2 * self.dimension
        i_n = np.arange(1, self.num_points // 2) # 索引n
        tem_s = np.sum(1 / np.power(i_n, tmp_1))
        tmp_2 = self.num_points ** 2 / (2 * tem_s)
        cof = np.sqrt(self.Rq ** 2 * tmp_2)  # 尺度系数C

        dft_cof = np.zeros(self.num_points, dtype=np.complex64)
        dft_cof[0] = 0
        dft_cof[self.num_points // 2] = dft_cof[0]

        for i_k in np.arange(1, self.num_points // 2):
            if not self.stable:
                r = cof * np.power(i_k, -tmp_1 / 2) * np.random.rand()
            else:
                r = cof * np.power(i_k, -tmp_1 / 2)
            p = np.random.uniform(0, 2*np.pi, 1)[0]
            dft_cof[i_k] = np.complex64(r * np.cos(p), r * np.sin(p))

        dft_cof_conj = np.conj(dft_cof[1 : self.num_points // 2])[:: -1]
        dft_cof[self.num_points // 2 + 1:] = dft_cof_conj
        self.altitude = np.real(scipy.fft.ifft(dft_cof))



if __name__ == '__main__':
    frac_p = FractalProfile(num_points=128, dimension=1.8, step=1, Rq=1, stable=True)
    frac_p.get_fractal_profile_DFT()
    frac_p.show_profile()
    print(frac_p.altitude)